'''
Author: goog
Date: 2022-01-10 14:33:17
LastEditTime: 2022-01-14 10:59:18
LastEditors: goog
Description: 
FilePath: /chengdu/TensorRT20220110/DetZM/cls_toyota.py
Time Limit Exceeded!
'''
import os
import cv2
import numpy as np
from PIL import Image

import torch
from torchvision import transforms
try:
    from model import resnet_56
except:
    from .model import resnet_56

class ToyotaCls(object):
    def __init__(self, weights, num_classes=3, device='cuda:0', imgsz=224):
        self.num_classes = num_classes
        self.device = torch.device(device)
        self.model = resnet_56(num_classes=self.num_classes)
        state_dict = torch.load(weights, map_location=self.device)
        self.model.load_state_dict(state_dict, strict=False)
        self.model.to(device)
        self.model.eval()
        self.imgsz = imgsz

        self.transform = transforms.Compose([
             transforms.Resize(256),
             transforms.CenterCrop(224),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                  std=[0.229, 0.224, 0.225])
        ])

    @torch.no_grad()
    def predict(self, source):
        #source = cv2.cvtColor(source, cv2.COLOR_BGR2RGB)
        source = Image.fromarray(np.uint8(source))
        source = self.transform(source)
        source = torch.unsqueeze(source, dim=0).to(self.device)
        result = self.model(source)
        result = torch.softmax(result, dim=1)
        label = result.argmax(1)
        result = result.squeeze().data.cpu().numpy()
        label = label.data.cpu().numpy()
        return result, label[0]
    
    def clsSHBColor(self, image1, image2, image3):
        """
        :param model: the color model of SHB
        :param image1, image2, image3: the local image under the YinTiao
        :color: 0: black, 1: brown, 2: white
        :return: the predicted color
        """
        colors = {'0':'black', '1':'brown', '2':'white'}
        pred1, label1 = self.predict(image1)
        pred2, label2 = self.predict(image2)
        pred3, label3 = self.predict(image3)
        
        pred = (pred1+pred2+pred3)/3.0
        index = int(np.argmax(pred))
        color = colors[str(index)]
        return index, color

    def clsZSTColor(self, image1, image2, image3):
        """
        :param model: the color model of ZST
        :param image1, image2, image3: the local image of ZST
        :return: the predicted color
        """
        colors = {'0':'WGW', '1':'FMW', '2':'ZMW', '3':'OMW'}
        pred1, label1 = self.predict(image1)
        pred2, label2 = self.predict(image2)
        pred3, label3 = self.predict(image3)
        # print(pred1)
        # print(pred2)
        # print(pred3)
        pred = (pred1+pred2+pred3)/3.0
        index = int(np.argmax(pred))
        color = colors[str(index)]
        return index, color

    def clsLBZColor(self, image1, image2, image3):
        """
        :param model: the color model of ZST
        :param image1, image2, image3: the local image of LaBaZhao
        :return: the predicted color
        """
        colors = {'0':'black', '1':'gray', '2':'other'}
        pred1, label1 = self.predict(image1)
        pred2, label2 = self.predict(image2)
        pred3, label3 = self.predict(image3)
        pred = (pred1+pred2+pred3)/3.0
        index = int(np.argmax(pred))
        color = colors[str(index)]
        return index, color
    def clsKK(self, roi):
        categorys = {'1': 'normal', '0': 'abnormal'}
        pred, label = self.predict(image1)
        index = int(np.argmax(pred))
        category = categorys[str(index)]
        return index, category
    def clsRD(self, roi):
        categorys = {'1': 'normal', '0': 'abnormal'}
        pred, label = self.predict(image1)
        index = int(np.argmax(pred))
        category = categorys[str(index)]
        return index, category
    def clsLS(self, roi):
        categorys = {'1': 'normal', '0': 'abnormal'}
        pred, label = self.predict(image1)
        index = int(np.argmax(pred))
        category = categorys[str(index)]
        return index, category

def readThrImages(srcPath):
    imgFiles = [x for x in os.listdir(srcPath) if x.endswith('.jpg')]
    number = np.random.randint(0, len(imgFiles)-1)
    image1 = cv2.imread(os.path.join(srcPath, imgFiles[number]))
    image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)
    number = np.random.randint(0, len(imgFiles)-1)
    image2 = cv2.imread(os.path.join(srcPath, imgFiles[number]))
    image2 = cv2.cvtColor(image2, cv2.COLOR_BGR2RGB)
    number = np.random.randint(0, len(imgFiles)-1)
    image3 = cv2.imread(os.path.join(srcPath, imgFiles[number]))
    image3 = cv2.cvtColor(image3, cv2.COLOR_BGR2RGB)

    return image1, image2, image3


if __name__ == '__main__':
    weights = '../CTB/color_0108/best_model_acc_95.88.pth'
    srcPath = '../shb/1'
    for i in range(10):
        image1, image2, image3 = readThrImages(srcPath=srcPath)
        shb = ToyotaCls(weights=weights, num_classes=3)
        index, color = shb.clsSHBColor(shb, image1, image2, image3)
        print('index: ', i, ', color:', color)
    print('over')